# Purpose: Main agent control loop.
# Changes:
# - Corrected Plan instantiation with explicit defaults/None for optional fields.

import logging
import asyncio
import json
from datetime import datetime, timezone # Ensure imported
from typing import Optional, List

# Use relative imports
# Make sure AgentExecutionError is imported
from ..core.exceptions import AgentError, ConfigError, LLMError, MCPClientError, AgentExecutionError
from .agent_types import Task, Plan, Step, AgentContext, MemoryRecord, AgentObservation
from .observer import Observer
from .planner import Planner
from .executor import Executor
from .reflector import Reflector
from .memory.short_term import ShortTermMemory
from .memory.base_memory import BaseMemory
from ..core.config import get_config_value
from ..mcp_client import initialize_mcp_client, close_mcp_client, get_mcp_session

# Use core logger setup
from ..core.logging_setup import get_logger
logger = get_logger(__name__) # Corrected logger acquisition


class Orchestrator:
    """Coordinates the Observe-Plan-Execute-Reflect loop for a task."""
    def __init__(self):
        try:
            self.observer = Observer()
            self.planner = Planner()
            self.executor = Executor()
            self.reflector = Reflector()
            self.memory: BaseMemory = ShortTermMemory(max_records=get_config_value("agent.memory.short_term_max_records", 200))
            self.max_steps = get_config_value("agent.max_steps", 15)
            self.max_consecutive_errors = get_config_value("agent.max_consecutive_errors", 3)
            logger.info("Orchestrator initialized successfully.")
        except (AgentError, ConfigError, LLMError) as e:
            logger.critical(f"Orchestrator initialization failed: {e}", exc_info=True)
            raise AgentError(f"Failed to initialize agent components: {e}") from e

    async def execute_task(self, task_description: str) -> Task:
        """Executes a task described in natural language."""
        # Correct Task instantiation with explicit None for optional fields
        task = Task(
            description=task_description,
            status="pending",
            initial_state=None,
            result=None,
            error_message=None
            # created_at/updated_at have defaults
        )
        agent_context = AgentContext(task=task)
        logger.info(f"--- Starting Task Execution: {task.id} ---")
        logger.info(f"Goal: {task.description}")

        try:
            await initialize_mcp_client()
            task.status = "running"
            logger.info(f"Task {task.id} status set to 'running'.")
        except (MCPClientError, ConfigError) as e:
            task.status = "failed"; task.error_message = f"MCP Connection Failed: {e}"
            logger.critical(f"Task {task.id} failed pre-run (MCP connection): {e}")
            return task # Return early

        step_count = 0; consecutive_errors = 0
        tool_descriptions: Optional[str] = None

        while step_count < self.max_steps and task.status == "running":
            step_count += 1
            internal_step_num = (agent_context.plan.steps[-1].id + 1) if agent_context.plan and agent_context.plan.steps else 1
            logger.info(f"\n--- Orchestrator: Loop Step {step_count}/{self.max_steps} (Plan Step: {internal_step_num}) ---")
            current_step_object: Optional[Step] = None

            try:
                # 1. Observe
                logger.debug("Orchestrator: Observing environment...")
                observation = await self.observer.get_observation(task_domain="web")
                agent_context.set_observation(observation)
                await self.memory.add_record(
                    MemoryRecord( type="observation", content=observation.model_dump(), related_task_id=task.id, importance=None)
                ) # Pass importance=None
                logger.info(f"Orchestrator: Observation received (Type: {observation.observation_type}).")
                # Check for critical observation errors
                if observation.observation_type == "web_page" and observation.data.error_message and "Critical" in observation.data.error_message:
                     raise AgentError(f"Critical observation error: {observation.data.error_message}")

                # 2. Plan
                logger.debug("Orchestrator: Planning next step...")
                if tool_descriptions is None:
                    tool_descriptions = await self.planner._load_tool_descriptions()

                next_step = await self.planner.decide_next_step(
                    task=agent_context.task,
                    current_observation=agent_context.current_observation, # Pass Optional
                    step_history=agent_context.step_history
                )
                if next_step is None: raise AgentError("Planner failed to return a next step.")
                current_step_object = next_step

                # Correct Plan instantiation with explicit defaults/None
                if agent_context.plan is None:
                    agent_context.plan = Plan(
                        task_id=task.id,
                        steps=[], # Start with empty steps
                        status="running", # Set status to running when created
                        original_prompt=None, # Explicitly None
                        raw_llm_response=None # Explicitly None
                        # id, created_at, updated_at have defaults
                    )
                agent_context.plan.steps.append(next_step)
                await self.memory.add_record(
                    MemoryRecord(type="plan_step", content=next_step.model_dump(), related_task_id=task.id, importance=None)
                ) # Pass importance=None

                # Check for finish_task before execution
                if next_step.tool_name == "finish_task":
                    reason = next_step.arguments.get('reason', 'Task finished by agent.')
                    logger.info(f"Orchestrator: Planner requested finish. Reason: {reason}")
                    task.status = "completed"; task.result = reason
                    next_step.status = "completed"; agent_context.add_executed_step(next_step)
                    break # Exit loop

                # 3. Execute
                logger.debug(f"Orchestrator: Executing step {next_step.id} (Tool: {next_step.tool_name})...")
                next_step.status = "running"; agent_context.add_executed_step(next_step)

                execution_result = await self.executor.execute_step(next_step)

                next_step.status = "completed"; next_step.result = execution_result
                agent_context.add_executed_step(next_step)
                await self.memory.add_record(
                    MemoryRecord(type="action_result", content={"step_id": next_step.id, "status": "completed", "result": repr(execution_result)[:500]}, related_task_id=task.id, importance=None)
                ) # Pass importance=None
                logger.info(f"Orchestrator: Step {next_step.id} execution successful.")
                consecutive_errors = 0 # Reset error count

            except (AgentError, MCPClientError, AgentExecutionError) as e: # Catch specific execution error too
                logger.error(f"Orchestrator: Error during loop step {step_count}: {e}", exc_info=False)
                consecutive_errors += 1
                step_to_update = current_step_object # Use the step that was being processed

                if step_to_update:
                    step_to_update.status = "failed"
                    step_to_update.error_message = str(e)
                    agent_context.add_executed_step(step_to_update) # Update history
                    await self.memory.add_record( MemoryRecord( type="error_result", content=step_to_update.model_dump(), related_task_id=task.id, importance=None))

                    # Reflect on the error
                    logger.warning("Orchestrator: Step failed. Initiating reflection...")
                    try:
                        if tool_descriptions is None: tool_descriptions = await self.planner._load_tool_descriptions()
                        reflection_insight = await self.reflector.reflect_on_step_error(
                            context=agent_context, failed_step=step_to_update, available_tools_desc=tool_descriptions
                        )
                        if reflection_insight:
                            logger.info(f"Reflection insight: {reflection_insight}")
                            await self.memory.add_record(MemoryRecord(type="reflection", content=reflection_insight, related_task_id=task.id, importance=None))
                    except Exception as reflect_err: logger.error(f"Error during reflection: {reflect_err}", exc_info=True)

                else: # Error occurred outside a specific step context (e.g., planning failure)
                    logger.error(f"Agent error not linked to a specific step: {e}")
                    # Maybe fail task immediately if planning fails?
                    # task.status = "failed"; task.error_message = str(e); break

                # Check for max consecutive errors
                if consecutive_errors >= self.max_consecutive_errors:
                    logger.error(f"Max consecutive errors ({self.max_consecutive_errors}) reached. Failing task.")
                    task.status = "failed"; task.error_message = f"Max errors reached. Last: {e}"
                    break

            except Exception as e: # Catch unexpected errors
                logger.critical(f"Orchestrator: Unexpected critical error in loop step {step_count}: {e}", exc_info=True)
                task.status = "failed"; task.error_message = f"Critical unexpected error: {e}"
                if current_step_object: # Try to mark current step as failed
                    current_step_object.status = "failed"; current_step_object.error_message = f"Critical error: {e}"
                    agent_context.add_executed_step(current_step_object)
                    await self.memory.add_record(MemoryRecord(type="error_result", content=current_step_object.model_dump(), related_task_id=task.id, importance=None))
                break # Exit loop

        # --- End of Loop ---
        if task.status == "running": # Reached max steps
            logger.warning(f"Task {task.id} reached max steps ({self.max_steps}).")
            task.status = "failed"; task.error_message = f"Max steps ({self.max_steps}) reached."

        task.updated_at = datetime.now(timezone.utc)
        logger.info(f"--- Task Execution Finished: {task.id} ---")
        logger.info(f"Final Status: {task.status.upper()}")
        if task.status == "failed": logger.error(f"Error: {task.error_message}")

        await self._cleanup_resources()
        return task

    async def _cleanup_resources(self):
        """Cleans up resources like the MCP client connection."""
        logger.debug("Orchestrator: Cleaning up MCP client session...")
        await close_mcp_client()

# --- Helper function ---
async def run_task_loop(initial_task_description: str):
    """Creates an orchestrator and runs a single task."""
    orchestrator: Optional[Orchestrator] = None
    try:
        orchestrator = Orchestrator()
        final_task_state = await orchestrator.execute_task(initial_task_description)
        logger.info("\n--- FINAL TASK STATE ---")
        logger.info(final_task_state.model_dump_json(indent=2))
        logger.info("------------------------")
    except Exception as e: # Catch errors during orchestrator init or top-level execution
        logger.critical(f"Agent task loop failed critically: {e}", exc_info=True)
    finally:
        logger.info("Ensuring MCP client cleanup after task loop attempt.")
        # Use direct call to close_mcp_client as orchestrator might be None
        await close_mcp_client()